package io.vertx.example.util;

import java.util.List;

import io.vertx.core.Handler;
import io.vertx.core.streams.Pump;
import io.vertx.core.streams.ReadStream;
import io.vertx.core.streams.WriteStream;
import io.vertx.core.streams.impl.PumpImpl;

public class MultiPumpImpl<T> implements Pump {


    private final ReadStream<T> readStream;
    private final WriteStream<T> writeStream;
    private final Handler<T> dataHandler;
    private final Handler<Void> drainHandler;
    private int pumped;



    public MultiPumpImpl(ReadStream<T> rs, List<WriteStream<T>> ws) {
        this.readStream = rs;
        this.writeStream = ws.get(0);
        drainHandler = v-> readStream.resume();
        dataHandler = data -> {
            // writeStream.write(data);
            for (WriteStream<T> item: ws) {
                item.write(data);
            }
            incPumped();
            if (writeStream.writeQueueFull()) {
                readStream.pause();
                writeStream.drainHandler(drainHandler);
            }
        };
    }

    /**
     * Set the write queue max size to {@code maxSize}
     */
    @Override
    public MultiPumpImpl setWriteQueueMaxSize(int maxSize) {
        writeStream.setWriteQueueMaxSize(maxSize);
        return this;
    }

    /**
     * Start the Pump. The Pump can be started and stopped multiple times.
     */
    @Override
    public MultiPumpImpl start() {
        readStream.handler(dataHandler);
        return this;
    }

    /**
     * Stop the Pump. The Pump can be started and stopped multiple times.
     */
    @Override
    public MultiPumpImpl stop() {
        writeStream.drainHandler(null);
        readStream.handler(null);
        return this;
    }

    /**
     * Return the total number of elements pumped by this pump.
     */
    @Override
    public synchronized int numberPumped() {
        return pumped;
    }

    // Note we synchronize as numberPumped can be called from a different thread however incPumped will always
    // be called from the same thread so we benefit from bias locked optimisation which should give a very low
    // overhead
    private synchronized void incPumped() {
        pumped++;
    }
}
